import dataclasses
from typing import TYPE_CHECKING

import flax.nnx as nnx
import jax
import jax.numpy as jnp
from typing_extensions import override

from openpi.models import model as _model
import openpi.models.gemma as _gemma
from openpi.shared import array_typing as at
import openpi.shared.nnx_utils as nnx_utils

if TYPE_CHECKING:
    from openpi.models.pi0 import Pi0


@dataclasses.dataclass(frozen=True)
class Pi0Config(_model.BaseModelConfig):
    dtype: str = "bfloat16"
    paligemma_variant: _gemma.Variant = "gemma_2b"
    action_expert_variant: _gemma.Variant = "gemma_300m"

    # Set the model specific defaults.
    action_dim: int = 32
    action_horizon: int = 50
    max_token_len: int = None  # type: ignore
    # Pi05 has two differences from Pi0:
    # - the state input is part of the discrete language tokens rather than a continuous input that is part of the suffix
    # - the action expert uses adaRMSNorm to inject the flow matching timestep
    pi05: bool = False
    # This config option is not used directly by the model, but it is read by the ModelTransformFactory.
    discrete_state_input: bool = None  # type: ignore

    # PRoPE configuration
    use_prope: bool = True
    head_dim: int = 64
    cameras: int = 3  # 根据实际相机数量调整
    image_width: int = 224
    image_height: int = 224

    def __post_init__(self):
        if self.max_token_len is None:
            object.__setattr__(self, "max_token_len", 200 if self.pi05 else 48)
        if self.discrete_state_input is None:
            object.__setattr__(self, "discrete_state_input", self.pi05)

    @property
    @override
    def model_type(self) -> _model.ModelType:
        if self.pi05:
            return _model.ModelType.PI05
        return _model.ModelType.PI0

    @override
    def create(self, rng: at.KeyArrayLike) -> "Pi0":
        from openpi.models.pi0 import Pi0

        return Pi0(self, rngs=nnx.Rngs(rng))

    @override
    def inputs_spec(self, *, batch_size: int = 1) -> tuple[_model.Observation, _model.Actions]:
        image_spec = jax.ShapeDtypeStruct([batch_size, *_model.IMAGE_RESOLUTION, 3], jnp.float32)
        image_mask_spec = jax.ShapeDtypeStruct([batch_size], jnp.bool_)

        with at.disable_typechecking():
            observation_spec = _model.Observation(
                images={
                    "base_0_rgb": image_spec,
                    "left_wrist_0_rgb": image_spec,
                    "right_wrist_0_rgb": image_spec,
                },
                image_masks={
                    "base_0_rgb": image_mask_spec,
                    "left_wrist_0_rgb": image_mask_spec,
                    "right_wrist_0_rgb": image_mask_spec,
                },
                state=jax.ShapeDtypeStruct([batch_size, self.action_dim], jnp.float32),
                tokenized_prompt=jax.ShapeDtypeStruct([batch_size, self.max_token_len], jnp.int32),
                tokenized_prompt_mask=jax.ShapeDtypeStruct([batch_size, self.max_token_len], bool),
            )
        action_spec = jax.ShapeDtypeStruct([batch_size, self.action_horizon, self.action_dim], jnp.float32)

        return observation_spec, action_spec

    # 添加自定义的fake_obs方法以支持PRoPE
    def fake_obs(self, batch_size: int = 1) -> _model.Observation:
        """Generate fake observation data with PRoPE support."""
        observation_spec, _ = self.inputs_spec(batch_size=batch_size)
        
        # 如果使用PRoPE，添加viewmats和Ks参数
        if self.use_prope:
            # 为每个相机创建视图矩阵(viewmats)和内参矩阵(Ks)
            camera_names = list(observation_spec.images.keys())
            
            # 初始化viewmats和Ks字典
            viewmats = {}
            Ks = {}
            
            # 为每个相机生成默认的viewmats和Ks
            for cam_name in camera_names:
                # viewmats: [batch, 4, 4] - 单位矩阵作为默认视图矩阵
                viewmats[cam_name] = jnp.tile(
                    jnp.eye(4, dtype=jnp.float32), [batch_size, 1, 1]
                )
                
                # Ks: [batch, 3, 3] - 默认内参矩阵
                K = jnp.array([
                    [1000.0, 0.0, 112.0],   # fx, 0, cx
                    [0.0, 1000.0, 112.0],   # 0, fy, cy
                    [0.0, 0.0, 1.0]         # 0, 0, 1
                ], dtype=jnp.float32)
                Ks[cam_name] = jnp.tile(K, [batch_size, 1, 1])
        else:
            viewmats = None
            Ks = None
        
        # 创建fake_observation时就包含所有字段
        fake_observation = _model.Observation(
            images=jax.tree.map(lambda x: jnp.ones(x.shape, x.dtype), observation_spec.images),
            image_masks=jax.tree.map(lambda x: jnp.ones(x.shape, x.dtype), observation_spec.image_masks),
            state=jnp.ones(observation_spec.state.shape, observation_spec.state.dtype),
            tokenized_prompt=jax.tree.map(lambda x: jnp.ones(x.shape, x.dtype), observation_spec.tokenized_prompt) if observation_spec.tokenized_prompt is not None else None,
            tokenized_prompt_mask=jax.tree.map(lambda x: jnp.ones(x.shape, x.dtype), observation_spec.tokenized_prompt_mask) if observation_spec.tokenized_prompt_mask is not None else None,
            viewmats=viewmats,
            Ks=Ks,
        )
        
        return fake_observation

    def get_freeze_filter(self) -> nnx.filterlib.Filter:
        """Returns the freeze filter based on the model config."""
        filters = []
        has_lora = False
        gemma_params_filter = nnx_utils.PathRegex(".*llm.*")
        action_expert_params_filter = nnx_utils.PathRegex(".*llm.*_1.*")
        if "lora" in self.paligemma_variant:
            filters.append(
                gemma_params_filter,
            )
            if "lora" not in self.action_expert_variant:
                # If only freeze gemma params, exclude action expert params.
                filters.append(
                    nnx.Not(action_expert_params_filter),
                )
            has_lora = True
        elif "lora" in self.action_expert_variant:
            filters.append(
                action_expert_params_filter,
            )
            has_lora = True

        if has_lora:
            # If any lora is used, exclude all lora params.
            filters.append(
                nnx.Not(nnx_utils.PathRegex(".*lora.*")),
            )
        if not filters:
            return nnx.Nothing
        return nnx.All(*filters)